//============================================================================
// Name        : CppFourierLawHB.cpp
// Author      : Carlos Olivares
// Version     : 0.5
// Copyright   : GPL
// Description : Molecular dynamics simulation
//============================================================================

#include <stdio.h>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <string>

#include <gsl/gsl_math.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_rng.h>
#include <time.h>

using namespace std;
#define DIM 3

int indx(int i, int j){
  return (i*DIM+j);
}

int indx(int i, int j, int Nrows){
  return (i*Nrows+j);
}

//DrPhiLJOverr
double DrPhiLJOverr(double delta_r2, double sigma2, double eps){
  //TODO: Make ri6, ri12 global variables.
  double ri6  = (sigma2/delta_r2)*(sigma2/delta_r2)*(sigma2/delta_r2);
  double ri12 = ri6*ri6;
  return -24*eps*(2*ri12 - ri6)/delta_r2;
}

// --------------------------------------------
double FSFAxis(double delta_r2, double deltaAxis, double sigma2, double eps, double rc, double DrPhiLJOverr_rc){
  if (delta_r2 < rc*rc){
    return - (-DrPhiLJOverr(delta_r2, sigma2, eps) + DrPhiLJOverr_rc )*deltaAxis;
  }else{
    return 0.0;
  }
}

double PhiLJ(double delta_r2, double sigma2, double eps){
  double ri6  = (sigma2/delta_r2)*(sigma2/delta_r2)*(sigma2/delta_r2);
  double ri12 = ri6*ri6;
  return 4*eps*(ri12 - ri6);
}

double PhiSF(double delta_r2, double sigma2, double eps, double rc, double PhiLJ_rc, double DrPhiLJOverr_rc){
  double rc2 = rc*rc;
  if (delta_r2 < rc2){
    return PhiLJ(delta_r2, sigma2, eps) - PhiLJ_rc - DrPhiLJOverr_rc*(delta_r2 - rc2)/2;
  }else{
    return 0.0;
  }
}

// Random Number generation with Maxwell distribution
double FOverG(double y){
 return y*exp(-0.5*y*y + y);
}

double randMaxwellDist(gsl_rng *r){
  double Y = 0.0;
  double U = 0.0;
  double c = 2.2039;
  double X = 0.0;

  do {
    Y = gsl_ran_exponential(r,1);
    U = gsl_rng_uniform(r);
  } while(U > FOverG(Y)/c);

  X = Y;
  return X;
}

// --------------------------------------------

class Gas{
  public:
    int N;
    double *vec_r;
    double *vec_v;
    double *mass;

    double xlen,ylen;
    double halfXlen, halfYlen;
    double sigma;
    
    double eps;
    double sigma2;
    double rc;
    double rc2;
    
    double PhiLJ_rc;
    double DrPhiLJOverr_rc;
    
    double *deltaX;
    double *deltaY;
    double *deltaZ;
    double *deltaR2;
    
    double *Fx;
    double *Fy;
    double *Fz;
    
    double *force;
    double *Ekin_n;
    double *halfV;
    
    double kinEner;
    double potEner;
        
    Gas(int const, double, double);
    ~Gas();
    void init1DVec_r(); // FIXME: (TEMPORAL)
    void init1DVec_v(double);
    void calculateDeltas();
    void calculateForces();
    void calculateHalfV();
    void calculateEkin_n();
    void calculatePotentialEnergy();
    void calculateKineticEnergy();
    void printVec_r();
    void fprint_vec_r(string);
    void fprint_vec_v(string);
    void readfile_vec_r(string);
    void readfile_vec_v(string);
    void printDeltaX();
    void printFx();
    void printForce();
};

Gas::Gas(int const N, double xlen, double ylen){
  this->N = N;
  this->vec_r = new double[N*DIM];
  this->vec_v = new double[N*DIM];
  this->xlen = xlen;
  this->ylen = ylen;
  this->halfXlen = 0.5*xlen;
  this->halfYlen = 0.5*ylen;
  this->sigma=1.0;
  this->mass = new double[N];
  
  this->deltaX = new double[N*N];
  this->deltaY = new double[N*N];
  this->deltaZ = new double[N*N];
  this->deltaR2= new double[N*N];
  
  this->Fx = new double[N*N];
  this->Fy = new double[N*N];
  this->Fz = new double[N*N];
    
  this->force = new double[N*DIM];
  
  this->Ekin_n = new double[N];
  this->halfV  = new double[N*N];
  
  this->kinEner =0.0;
  this->potEner=0.0;
  
  this->eps   = 0.1;
  this->sigma2= this->sigma*this->sigma;
  this->rc    = 2.5;
  this->rc2   = this->rc*this->rc;

  this->PhiLJ_rc   = PhiLJ(this->rc2, this->sigma2, this->eps);
  this->DrPhiLJOverr_rc = DrPhiLJOverr(this->rc2, this->sigma2, this->eps);
  
  // Initiate vectors Row-major array
  for (int i=0; i<this->N; i++){
    for (int j=0; j<DIM; j++){
      this->vec_r[indx(i,j)] = 0.0;
      this->vec_v[indx(i,j)] = 0.0;
      this->force[indx(i,j)] = 0.0;
    }
  }
  
  for (int l=0; l < this->N; l++){
    for (int k=0; k < this->N; k++){
      this->Fx[indx(l,k,this->N)] = 0.0;
      this->Fy[indx(l,k,this->N)] = 0.0;
      this->Fz[indx(l,k,this->N)] = 0.0;
    }
  }
  for (int i=0; i< this->N; i++){
    this->mass[i] =1.0;
  }
  
  for (int i=0; i< this->N-1;i++){
    for (int j=0; j<DIM; j++){
      this->vec_r[indx(i,j)] = 0.0;
      this->vec_v[indx(i,j)] = 0.0;
    }
  }
  
  int indx2=0;
  for (int i=0; i < this->N; i++){
    for (int j=0; j < this->N; j++){
      indx2 = i*this->N+j;
      this->deltaX[indx2] = 0.0;
      this->deltaY[indx2] = 0.0;
      this->deltaZ[indx2] = 0.0;
      this->deltaR2[indx2]= 0.0;
    }
  }
  
};

// FIXME: (TEMPORAL) 
void Gas::init1DVec_r(){
  
  double delta_r = 3.0;//(this->xlen)/(this->N+1);
  int j=0; // just for x axis, y=z=0.0
  if (delta_r > this->sigma){
    int h=this->N/2;
    if (this->N%2==1){
      this->vec_r[indx(0,j)] = -h*delta_r;
    }else{
      this->vec_r[indx(0,j)] = -h*delta_r + delta_r/2.0;
    }  
    for (int i=0; i< this->N-1;i++){
      this->vec_r[indx(i+1,j)] = this->vec_r[indx(i,j)]+delta_r;
    }
  }else{
    cout<<"Density to high, decrease to avoid overlap"<<endl;
  }
};

// FIXME: All the initial velocities are Zero (TEMPORAL)
void Gas::init1DVec_v(double Temp){
    
  const gsl_rng_type *T; T = gsl_rng_default;
  gsl_rng *r; r = gsl_rng_alloc(T);
  long unsigned int seed;
  seed = 100;//(long unsigned int) time(NULL); // Get the time of the system as seed
  gsl_rng_set(r,seed);
  
  int j=0; // just for vx , vy=vz=0.0
  double velmean=0.0;
  for (int i=0; i<this->N; i++){
    this->vec_v[indx(i,j)] = sqrt(Temp)*gsl_ran_gaussian(r, 1);
    velmean+= this->vec_v[indx(i,j)];
  }
  velmean=velmean/N;
  for (int i=0; i<this->N; i++){
    this->vec_v[indx(i,j)] = this->vec_v[indx(i,j)]-velmean;
  }
  
  this->calculateEkin_n();
  this->calculateKineticEnergy();
  double TempFactor=sqrt(0.5*this->N*Temp/this->kinEner);
  
  cout<<"kinEner="<<kinEner<<endl;
  cout<<"TempFactor="<<TempFactor<<endl;
  for (int i=0; i<this->N; i++){
    this->vec_v[indx(i,j)] = TempFactor*this->vec_v[indx(i,j)];
  }
  this->calculateEkin_n();
  this->calculateKineticEnergy();
  cout<<"kinEnerScaled="<<kinEner<<endl;
  gsl_rng_free (r);
};


void Gas::calculateDeltas(){
  int indx2=0;
  for (int i=0; i < this->N; i++){
    for (int j=0; j < this->N; j++){
      indx2 = i*this->N+j;
      this->deltaX[indx2] = this->vec_r[indx(j,0)] - this->vec_r[indx(i,0)];
      this->deltaY[indx2] = this->vec_r[indx(j,1)] - this->vec_r[indx(i,1)];
      this->deltaZ[indx2] = this->vec_r[indx(j,2)] - this->vec_r[indx(i,2)];
      this->deltaR2[indx2]= this->deltaX[indx2]*this->deltaX[indx2]
                          + this->deltaY[indx2]*this->deltaY[indx2]
                          + this->deltaZ[indx2]*this->deltaZ[indx2];
    }
  }
}    

void Gas::calculateForces(){
  for (int l=0; l < this->N; l++){
    for (int k=0; k < l; k++){
      this->Fx[indx(l,k,this->N)] = FSFAxis(this->deltaR2[indx(l,k,this->N)], this->deltaX[indx(l,k,this->N)], 
              this->sigma2, this->eps, this->rc, this->DrPhiLJOverr_rc);
      this->Fy[indx(l,k,this->N)] = FSFAxis(this->deltaR2[indx(l,k,this->N)], this->deltaY[indx(l,k,this->N)], 
              this->sigma2, this->eps, this->rc, this->DrPhiLJOverr_rc);
      this->Fz[indx(l,k,this->N)] = FSFAxis(this->deltaR2[indx(l,k,this->N)], this->deltaZ[indx(l,k,this->N)], 
              this->sigma2, this->eps, this->rc, this->DrPhiLJOverr_rc);
      
      this->Fx[indx(k,l,this->N)] = -this->Fx[indx(l,k,this->N)];
      this->Fy[indx(k,l,this->N)] = -this->Fy[indx(l,k,this->N)];
      this->Fz[indx(k,l,this->N)] = -this->Fz[indx(l,k,this->N)];
    }
  }
  
  for (int i=0; i< this->N; i++){
    this->force[indx(i,0)] = 0.0;
    this->force[indx(i,1)] = 0.0;
    this->force[indx(i,2)] = 0.0;
    for (int k=0; k< this->N; k++){
      this->force[indx(i,0)] +=  this->Fx[indx(i,k,this->N)];
      this->force[indx(i,1)] +=  this->Fy[indx(i,k,this->N)];
      this->force[indx(i,2)] +=  this->Fz[indx(i,k,this->N)];
    }
  }
  
}

void Gas::calculateHalfV(){
  for(int l=0; l<this->N; l++){
    for (int k=0; k < l; k++){
      this->halfV[indx(l,k,this->N)] = 0.5*PhiSF(this->deltaR2[indx(l,k,this->N)], this->sigma2, 
              this->eps, this->rc, this->PhiLJ_rc, this->DrPhiLJOverr_rc);
      this->halfV[indx(k,l,this->N)] = this->halfV[indx(l,k,this->N)];
    }
  }
}
void Gas::calculateEkin_n(){
  for(int i=0; i< this->N; i++){
    this->Ekin_n[i]= 0.0;
    for(int j=0; j<DIM; j++){
      this->Ekin_n[i]+=this->vec_v[indx(i,j)]*this->vec_v[indx(i,j)];
    }
    this->Ekin_n[i]=0.5*mass[i]*this->Ekin_n[i];
  }
}

void Gas::calculatePotentialEnergy(){
  this->potEner=0.0;
  for(int l=0; l<this->N; l++){
    for (int k=0; k < this->N; k++){
      this->potEner += this->halfV[indx(l,k,this->N)];
    }
  }  
}

void Gas::calculateKineticEnergy(){
  this->kinEner = 0.0;
  for (int i=0; i< this->N; i++){
    this->kinEner += this->Ekin_n[i];
  }
}

Gas::~Gas(){
  delete[] vec_r;
  delete[] vec_v;
  delete[] deltaX;
  delete[] deltaY;
  delete[] deltaZ;
  delete[] deltaR2;
  delete[] Fx;
  delete[] Fy;
  delete[] Fz;
  delete[] mass;
  delete[] force;
  delete[] Ekin_n;
  delete[] halfV;
};

void Gas::printVec_r(){
  
  cout<<"__________"<<endl;
  cout<<"vec_r"<<endl;
  for (int i=0; i<this->N; i++){
    for (int j=0; j<DIM; j++){
      cout<< this->vec_r[indx(i,j)]<<" ";
    }
    cout<<endl;
  }
  cout<<"__________"<<endl;
};

void Gas::fprint_vec_r(string outfilename){

  ofstream ofile;
  ofile.open(outfilename.c_str());
  for (int i=0; i<this->N; i++){
    for (int j=0; j<DIM; j++){
      ofile << this->vec_r[indx(i,j)]<<" ";
    }
    ofile << endl;
  }
  ofile.close();
};

void Gas::fprint_vec_v(string outfilename){

  ofstream ofile;
  ofile.open(outfilename.c_str());
  for (int i=0; i<this->N; i++){
    for (int j=0; j<DIM; j++){
      ofile << this->vec_v[indx(i,j)]<<" ";
    }
    ofile << endl;
  }
  ofile.close();
};

void Gas::readfile_vec_r(string inputfilename){
  ifstream ifile;
  string line;
  ifile.open(inputfilename.c_str());
  
  for (int i=0; i<this->N; i++){
    for (int j=0; j<DIM; j++){
      ifile >> this->vec_r[indx(i,j)];
    }
  }
  ifile.close();
};

void Gas::readfile_vec_v(string inputfilename){
  ifstream ifile;
  string line;
  ifile.open(inputfilename.c_str());
  
  for (int i=0; i<this->N; i++){
    for (int j=0; j<DIM; j++){
      ifile >> this->vec_v[indx(i,j)];
    }
  }
  ifile.close();
};

void Gas::printDeltaX(){
for (int l=0; l< this->N; l++){
    for (int k=0; k< this->N; k++){
      cout<< this->deltaX[indx(l, k, this->N)]<<" ";
    }
    cout<<endl;
  }
}

void Gas::printFx(){
for (int l=0; l< this->N; l++){
    for (int k=0; k< this->N; k++){
      cout<< this->Fx[indx(l, k, this->N)]<<" ";
    }
    cout<<endl;
  }
}

void Gas::printForce(){
for (int i=0; i< this->N; i++){
    for (int j=0; j< DIM; j++){
      cout<< this->force[indx(i,j)]<<" ";
    }
    cout<<endl;
  }
}

/**************************************/
/**************************************/
/**************************************/
/**************************************/
/**************************************/
/**************************************/
int main(int argc, char* argv[]) {
  double const gasDensity=3.0;
  int const N = atoi(argv[1]); //40
  int const Ni = 1./N;
  double xlen=gasDensity*(N+1), ylen=4;
  double T1=atof(argv[2]);//4.0;
  double T2=atof(argv[3]);//10.0;
  double dm=atof(argv[4]); // 0.2;
  double otherDeltaR=atof(argv[5]);
//  string strIniFilename_vec_r(argv[6]);;
//  string strIniFilename_vec_v(argv[7]);
  
  double *TempProf = new double[N];
  
  string strN(argv[1]);
  string strT1(argv[2]);
  string strT2(argv[3]);
  string strdm(argv[4]);
  string strODR(argv[5]);
  
  
//  ofstream aux1file;;
//  ofstream aux2file;
//  aux1file.open("left0.dat");
//  aux2file.open("right0.dat");
  
  
  string outfilename ="N_"+strN+"__T1_"+strT1+"__T2_"+strT2+"__dm_"+strdm+"__dr_"+strODR+".dat";
  string outfilename2="N_"+strN+"__T1_"+strT1+"__T2_"+strT2+"__dm_"+strdm+"__dr_"+strODR+".trv";
  cout<<outfilename<<endl;
  ofstream outfile;
  ofstream outfile2;
  
  outfile.open(outfilename.c_str());
  outfile2.open(outfilename2.c_str());
  // FILE header
  outfile2 << "#ttime ";
  for (int i=0; i<N;i++){
    outfile2 << "x_"<< i <<" ";
  }
  for (int i=0; i<N;i++){
    outfile2 << "vx_"<< i <<" ";
  }
  outfile2 << endl;
  
  
  // Initial conditions
  Gas gas_old(N, xlen, ylen); Gas gas_new(N, xlen, ylen);
 
    // Mass distribution
  for (int i=0; i<gas_new.N; i++){
    if (i%2==0){
      gas_old.mass[i] =1.0;
      gas_new.mass[i] =1.0;
    }else{
      gas_old.mass[i] =1.0+dm;
      gas_new.mass[i] =1.0+dm;
    }
  }
  
  gas_old.init1DVec_r(); gas_old.init1DVec_v(T1);
  gas_new.init1DVec_r(); gas_new.init1DVec_v(T1);
  for (int i=0; i<N; i++){
    gas_new.vec_v[indx(i,0)]= gas_old.vec_v[indx(i,0)];
  }
  
//  // TODO: Initial conditions to test the energy
//  if (argc>6){
//    gas_old.readfile_vec_r(strIniFilename_vec_r);
//    gas_old.readfile_vec_v(strIniFilename_vec_v);
//
//    gas_old.readfile_vec_r(strIniFilename_vec_r);
//    gas_old.readfile_vec_v(strIniFilename_vec_v);
//  }
  
 
  
  for (int i=0; i<gas_new.N; i++){
    cout<<gas_new.mass[i];
    cout<<" ";
  }
  
  gas_old.calculateDeltas();
  gas_old.calculateForces();  

  const gsl_rng_type *T; T = gsl_rng_default;
  gsl_rng *r; r = gsl_rng_alloc(T);
  long unsigned int seed;
  seed = 100;//(long unsigned int) time(NULL); // Get the time of the system as seed
  gsl_rng_set(r,seed);
  
  double dt=0.0001;
  double dt2 = dt*dt;  

  // TODO Aqui comienza el While!!!
  int ttime = 0;
  int total_time = 40000000;
  int timeDim = 400000;
  int taux = total_time/timeDim;
  int other_time = 10000000;
  
  while (ttime < total_time){

    // Position update.
    for (int i=0; i< gas_old.N; i++){
      for (int j=0; j< DIM; j++){
        gas_new.vec_r[indx(i,j)] = gas_old.vec_r[indx(i,j)]+dt*gas_old.vec_v[indx(i,j)]
            + 0.5*dt2*gas_old.force[indx(i,j)]/gas_old.mass[i];
      }
    }

    gas_new.calculateDeltas();
    gas_new.calculateForces();
    
    double deltaAux=0.0;
    // Velocity update.
    for (int i=0; i< gas_old.N; i++){
      for (int j=0; j< DIM; j++){
        gas_new.vec_v[indx(i,j)] = gas_old.vec_v[indx(i,j)]
            + 0.5*dt*(gas_old.force[indx(i,j)] + gas_new.force[indx(i,j)])/gas_new.mass[i];
      }
      // TODO: Validar si esta en la caja!!! y el HEAT BATH
      if (gas_new.vec_r[indx(i,0)] < -gas_new.halfXlen ){
        //cout<<"out of box left "<< ttime <<" "<< gas_new.vec_r[indx(i,0)]<<endl;
        
        // TODO: correct the positions.
        deltaAux= -gas_new.vec_r[indx(i,0)] - gas_new.halfXlen;
        gas_new.vec_r[indx(i,0)] = -gas_new.halfXlen + deltaAux;
        
        // new velocity FIXME: (TEMPORAL) invert velocity
        if (ttime > other_time)
          gas_new.vec_v[indx(i,0)] = sqrt(T1)*randMaxwellDist(r);
        else
          gas_new.vec_v[indx(i,0)] = -gas_new.vec_v[indx(i,0)];
        gas_new.vec_v[indx(i,1)] = 0.0 ;
        gas_new.vec_v[indx(i,2)] = 0.0 ;
        //aux1file << dt*ttime <<" "<< gas_new.vec_v[indx(i,0)] << endl;
      }else if (gas_new.vec_r[indx(i,0)] > gas_new.halfXlen){
        //cout<<"out of box right "<< ttime <<" "<< gas_new.vec_r[indx(i,0)]<<endl;
        deltaAux= gas_new.vec_r[indx(i,0)] - gas_new.halfXlen;
        gas_new.vec_r[indx(i,0)] = gas_new.halfXlen - deltaAux;
        
        // new velocity FIXME: (TEMPORAL) invert velocity
        if (ttime>other_time)
          gas_new.vec_v[indx(i,0)] = -sqrt(T2)*randMaxwellDist(r); //fabs(gsl_ran_gaussian(r,1));
        else
          gas_new.vec_v[indx(i,0)] = -gas_new.vec_v[indx(i,0)];
        gas_new.vec_v[indx(i,1)] = 0.0 ;
        gas_new.vec_v[indx(i,2)] = 0.0 ;
        //aux2file << dt*ttime <<" "<< gas_new.vec_v[indx(i,0)] << endl;
      }
    }

    // Save data in File.
    if (ttime % taux == 0){
      outfile << dt*ttime << " ";
//      outfile2 << dt*ttime << " ";
//      // FIXME: (TEMPORAL)
//      for (int i=0; i<gas_new.N; i++){
//        outfile2 << gas_new.vec_v[indx(i,0)] << " ";
//        outfile2 << gas_new.vec_r[indx(i,0)] << " ";
//      }
//      outfile2<<endl;
      gas_new.calculateEkin_n();
      gas_new.calculateHalfV();
      gas_new.calculateKineticEnergy();
      gas_new.calculatePotentialEnergy();
      outfile << gas_new.kinEner << " ";
      outfile << gas_new.potEner << " ";
      outfile << endl;
    }
    
    for (int i=0; i<gas_old.N; i++){
      for (int j=0; j<DIM; j++){
        gas_old.vec_r[indx(i,j)] = gas_new.vec_r[indx(i,j)];
        gas_old.vec_v[indx(i,j)] = gas_new.vec_v[indx(i,j)];
        gas_old.force[indx(i,j)] = gas_new.force[indx(i,j)];
      }
    }
    ttime++;
  }
  cout<< "#+++++++++++++++++++" <<endl;
  outfile.close();
  outfile2.close();
  gas_old.printVec_r();
  gsl_rng_free (r);
  return 0;
}

